﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Diagnostics;
using alib.Enumerable;
using alib.Debugging;

namespace agree
{
	using mrs;

	public class DmrsExtractor : m_observable<Mrs, DependencyMrs>
	{

		public DmrsExtractor(IRtParent parent, Mrs item)
			: base(parent, item)
		{
		}
		protected override void start(Mrs item)
		{
			//TokenSet ts = t0.FindTrace<TokLoader>().TokenSet;

			var dmrs = new DependencyMrs(item);
			_item(dmrs);

			_remove_task();
		}
	};


	public class DependencyMrs : identity_base
	{
		public string[] preds;
		public bool[] is_quantifier;

		public List<DMRSLink>[] links;
		public List<DMRSLink>[] rev_links;
		public int root_index;
		public List<HashSet<int>> lbl_eq;

		public Mrs mrs;

		public enum DMRSLinkType
		{
			NONE,
			QEQ,
			EQ,
			NEQ
		};

		//[DebuggerBrowsable(DebuggerBrowsableState.RootHidden)]
		//public String[] _for_dbg { get { return preds.Select(p => g.tm.TypeNameOrStringValue(p)).ToArray(); } }

		public class DMRSLink
		{
			public DMRSLink(DMRSLinkType link_type, string idx_arg, DMRSLinkType lbl_link_type, string lbl_arg, int from, int to)
			{
				this.link_type = link_type;
				this.idx_arg = idx_arg;
				this.lbl_link_type = lbl_link_type;
				this.lbl_arg = lbl_arg;
				this.to = to;
				this.from = from;
			}

			public DMRSLinkType link_type;
			public string idx_arg;
			public string lbl_arg;
			public DMRSLinkType lbl_link_type;
			public int to;
			public int from;
		}

		protected class TempLink
		{
			public DMRSLinkType lt;
			public string arg;

			public TempLink(DMRSLinkType lt, string arg)
			{
				this.lt = lt;
				this.arg = arg;
			}
		}

		public enum DmrsTreenessResult
		{
			NOT_A_TREE,
			TREE_EXCEPT_LABELS,
			TREE,
			TREE_AFTER_FILTERING
		}

		public DmrsTreenessResult CheckTreeness(HashSet<string> ignore_preds)
		{
			bool[] ignoreRels = preds.Select(p => ignore_preds.Contains(p)).ToArray();
			var num_links = links.Flatten<DMRSLink>().Count(l => !ignoreRels[l.to] && !ignoreRels[l.from]);
			var num_preds = ignoreRels.Count(b => !b);

			if (num_links != num_preds - 1)
				return DmrsTreenessResult.NOT_A_TREE;

			if (lbl_eq.Count > 0)
				return DmrsTreenessResult.TREE_EXCEPT_LABELS;

			if (links.Flatten<DMRSLink>().Count() == preds.Length - 1)
				return DmrsTreenessResult.TREE;

			return DmrsTreenessResult.TREE_AFTER_FILTERING;
		}


		public DependencyMrs(Mrs mrs)
			: base(mrs)
		{
			this.mrs = mrs;

			int c = mrs.RELS.Count;

			var arg0s = new Variable[c];
			var lbls = new Variable[c];

			this.preds = new string[c];
			//this.token_spans = new ChartSpan[c];
			this.is_quantifier = new bool[c];

			//DMRSLink objects in both forward and reverse directions
			//forward links have ARG1-3 pointing to ARG0
			this.links = new List<DMRSLink>[c];
			this.rev_links = new List<DMRSLink>[c];

			//sets of eps that share labels
			this.lbl_eq = new List<HashSet<int>>();

			TempLink[,] temp_links = new TempLink[c, c];
			TempLink[,] temp_lbl_links = new TempLink[c, c];

			var main_rel = new Dictionary<Variable, int>();
			var label_eq = new Dictionary<Variable, HashSet<int>>();
			var lbl_heads = new Dictionary<Variable, int>();
			var qeqs = new Dictionary<Variable, Variable>();

			//first pass through all rels:
			//1. store token span and pred for each EP (by index in the mrs object)
			//2. identify the non-quant EP that has each variable as its ARG0 and store in dictionary
			//3. build up a dictionary of LBL->EPs with that LBL
			for (int i = 0; i < c; i++)
			{
				var rel = mrs.RELS[i];

				//if (ts != null)
				//	token_spans[i] = ts.CharacterSpanToMinimalSpan(mrs.rels[i].CharacterSpan);
				links[i] = new List<DMRSLink>();
				rev_links[i] = new List<DMRSLink>();
				preds[i] = rel.PRED;

				//identify EP with ARG0 for each var
				//arg0_slots[i] = rel_slots[i][mrsm.ifeat_arg0];

				arg0s[i] = (Variable)rel["ARG0"];
				var rstr_var = rel["RSTR"];

				is_quantifier[i] = arg0s[i] != null && rstr_var != null;
				if (arg0s[i] != null && !is_quantifier[i])
				{
					main_rel[(Variable)rel["ARG0"]] = i;
				}

				lbls[i] = (Variable)rel["LBL"];
				if (lbls[i] != null)
				{
					if (!label_eq.ContainsKey(lbls[i]))
						label_eq[lbls[i]] = new HashSet<int>();
					label_eq[lbls[i]].Add(i);
				}
			}


			//Second pass through EPs
			//1. add links for quantifiers
			//2. if a normal EP shares its ARG0 with a quantifier, that EP is head of its LBL group
			for (int i = 0; i < c; i++)
			{
				if (is_quantifier[i])
				{
					if (arg0s[i] != null && main_rel.ContainsKey(arg0s[i]))
					{
						int head_idx = main_rel[arg0s[i]];
						lbl_heads[lbls[i]] = head_idx;

						var link = new DMRSLink(DMRSLinkType.NEQ, "ARG0", DMRSLinkType.QEQ, "RSTR", i, head_idx);
						links[i].Add(link);
						rev_links[head_idx].Add(link);
					}
				}
			}

			//loop through hcons to find all qeq relations
			foreach (var qeq in mrs.HCONS)
			{
				qeqs[qeq.HARG] = qeq.LARG;
			}

			//identify a "head" for each group of rels with the same label
			//also identify rels that share labels but are not connected by arguments 
			//(i.e. "The dog whose toy the cat bit barked")

			//for groups with > 1 EP, we can eliminate candidates for head of the group if they have an argument pointing to
			//ARG0 of another member of the group.
			foreach (var lbl_set in label_eq)
			{
				if (lbl_set.Value.Count > 1)
				{
					HashSet<int> working_set = new HashSet<int>(lbl_set.Value);

					//for each EP in the lbl_set, look for other EPs in the same set with arguments pointing to it
					foreach (int idx in lbl_set.Value)
					{
						if (arg0s[idx] == null)
							continue;

						foreach (int idx2 in lbl_set.Value.Where(x => x != idx))
						{
							foreach (var role in mrs.RELS[idx2].Where(role =>
								role.name != "ARG0" &&
								role.var == arg0s[idx]))
							{
								if (!lbl_heads.ContainsKey(lbl_set.Key) &&
									role.name != "LBL" &&
									role.name != "RSTR")
								{
									lbl_heads[lbl_set.Key] = idx;
								}

								working_set.Remove(idx);
							}
						}
					}

					if (working_set.Count > 1)
					{
						lbl_eq.Add(working_set);
					}
				}

				//if a label equiv set has only EP, make that EP head
				else
				{
					lbl_heads[lbl_set.Key] = lbl_set.Value.First();
				}
			}

			//for (RelIx ri = arg0.first; ri./*Next.*/rel != null; ri = ri.Next)	// ? (gcs)
			//{
			//	Relation cur = ri.rel as Relation;

			//	//I think qeqs might be visited, but they are not real EPs so skip
			//	//also, no reason to look at the EP we are already looking at
			//	if (cur == null || ri.i_feat == mrs.Arg0Feat /*or perhaps: || cur == rel ? */)
			//		continue;

			//	if (lbl_set.Value.Contains(cur.index))
			//	{
			//		if (!rel.FeatProps.Any(fp =>
			//		{
			//			Variable v = fp.p as Variable;
			//			return v != null &&
			//			fp.i_feat != mrsm.ifeat_arg0 &&
			//			fp.i_feat != mrsm.RelLblPath.LastIx &&
			//			fp.i_feat != mrsm.ifeat_rstr &&
			//			main_rel.ContainsKey(v) &&
			//			lbl_set.Value.Contains(main_rel[v]);
			//		}))
			//		{
			//			if (!lbl_heads.ContainsKey(lbl_set.Key))
			//			{
			//				lbl_heads[lbl_set.Key] = idx;
			//			}
			//		}

			//		working_set.Remove(cur.index);
			//	}
			//}
			//}

			//if more than one EP remains after this heuristic, we need to store the remaining items in working set
			//as having the same label, but not sharing any other argument positions.

			//collect "TempLinks" information about label and index connections between rels
			for (int from_idx = 0; from_idx < mrs.RELS.Count; from_idx++)
			{
				foreach (var role in mrs.RELS[from_idx])
				{
					if (role.name == "ARG0" || role.name == "LBL" || role.name == "RSTR")
						continue;

					if (main_rel.ContainsKey(role.var))
					{
						int to_idx = main_rel[role.var];

						DMRSLinkType eq_neq = DMRSLinkType.NEQ;
						if (lbls[from_idx] == lbls[to_idx])
							eq_neq = DMRSLinkType.EQ;

						temp_links[from_idx, to_idx] = new TempLink(eq_neq, role.name);
					}
					else if (lbl_heads.ContainsKey(role.var))
					{
						int to_idx = lbl_heads[role.var];
						temp_lbl_links[from_idx, to_idx] = new TempLink(DMRSLinkType.EQ, role.name);
					}
					else if (qeqs.ContainsKey(role.var) && lbl_heads.ContainsKey(qeqs[role.var]))
					{
						int to_idx = lbl_heads[qeqs[role.var]];
						temp_lbl_links[from_idx, to_idx] = new TempLink(DMRSLinkType.QEQ, role.name);
					}
					else if (role.var is Handle)
					{
						var link = new DMRSLink(DMRSLinkType.NONE, null, DMRSLinkType.EQ, role.name, from_idx, -1);
						links[from_idx].Add(link);
					}
					else
					{
						var link = new DMRSLink(DMRSLinkType.NEQ, role.name, DMRSLinkType.NONE, null, from_idx, -1);
						links[from_idx].Add(link);
					}
				}
			}

			//combine the collected information into full "DMRSLink" objects
			for (int i = 0; i < c; i++)
			{
				for (int j = 0; j < c; j++)
				{
					if (i == j || (temp_links[i, j] == null && temp_lbl_links[i, j] == null))
						continue;

					DMRSLinkType main_lt = DMRSLinkType.NONE, lbl_lt = DMRSLinkType.NONE;
					string idx_arg = null, lbl_arg = null;

					if (temp_links[i, j] != null)
					{
						main_lt = temp_links[i, j].lt;
						idx_arg = temp_links[i, j].arg;
					}
					if (temp_lbl_links[i, j] != null)
					{
						lbl_lt = temp_lbl_links[i, j].lt;
						lbl_arg = temp_lbl_links[i, j].arg;
					}

					DMRSLink new_link = new DMRSLink(main_lt, idx_arg, lbl_lt, lbl_arg, i, j);
					links[i].Add(new_link);
					rev_links[j].Add(new_link);
				}
			}

			var h = mrs.HOOK;
			if (h.LTOP != null && lbl_heads.ContainsKey(h.LTOP))
				root_index = lbl_heads[h.LTOP];
			else if (h.INDEX != null && main_rel.ContainsKey(h.INDEX))
				root_index = main_rel[h.INDEX];
			else root_index = -1;
		}
	}
}
